package tacGen;

import ic.ast.BinaryOp;
import ic.ast.BinaryOp.BinaryOps;
import ic.ast.Call;
import ic.ast.ClassType;
import ic.ast.DeclClass;
import ic.ast.DeclField;
import ic.ast.DeclLibraryMethod;
import ic.ast.DeclMethod;
import ic.ast.DeclStaticMethod;
import ic.ast.DeclVirtualMethod;
import ic.ast.HasObject;
import ic.ast.Length;
import ic.ast.Literal;
import ic.ast.LocalVariable;
import ic.ast.NewArray;
import ic.ast.NewInstance;
import ic.ast.Node;
import ic.ast.Parameter;
import ic.ast.PrimitiveType;
import ic.ast.PrimitiveType.DataType;
import ic.ast.Program;
import ic.ast.RefArrayElement;
import ic.ast.RefField;
import ic.ast.RefVariable;
import ic.ast.Statement;
import ic.ast.StaticCall;
import ic.ast.StmtAssignment;
import ic.ast.StmtBlock;
import ic.ast.StmtBreak;
import ic.ast.StmtCall;
import ic.ast.StmtContinue;
import ic.ast.StmtIf;
import ic.ast.StmtReturn;
import ic.ast.StmtWhile;
import ic.ast.This;
import ic.ast.UnaryOp;
import ic.ast.VirtualCall;
import ic.ast.Visitor;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Stack;

import sun.reflect.generics.reflectiveObjects.NotImplementedException;

public class TacVisitor implements Visitor {

	private static final String MAIN_LABEL = ":main";

	private int next_register = 0;
	private final Stack<Integer> context_stack = new Stack<Integer>();

	private final StringBuffer sb = new StringBuffer();

	private final StringBuffer sb_data = new StringBuffer();

	private String curr_class;

	private static final boolean flgTranslateLibraryCalls = true;

	private static final String error_msg_label_prefix = ":error_msg_";

	private static final String error_func_label_prefix = ":error_func_";

	private static final String null_deref_postfix = "null_deref";

	private static final String lblMsgNullDereference = error_msg_label_prefix
			+ null_deref_postfix;

	private static final String lblFuncNullDereference = error_func_label_prefix
			+ null_deref_postfix;

	private static final String array_bounds_postfix = "array_bounds";

	private static final String lblMsgArrayOutOfBounds = error_msg_label_prefix
			+ array_bounds_postfix;

	private static final String lblFuncArrayOutOfBounds = error_func_label_prefix
			+ array_bounds_postfix;

	private static final String array_size_postfix = "array_size";

	private static final String lblMsgNegativeArraySize = error_msg_label_prefix
			+ array_size_postfix;

	private static final String lblFuncNegativeArraySize = error_func_label_prefix
			+ array_size_postfix;

	private static final String div_zero_postfix = "div_zero";

	private static final String lblMsgDivByZero = error_msg_label_prefix
			+ div_zero_postfix;

	private static final String lblFuncDivByZero = error_func_label_prefix
			+ div_zero_postfix;

	private static final String lblErrorThenExit = ":error_then_exit";

	class VirtualMethodDetails {
		String method_name;
		String class_name;
		public int index;
	}

	class ClassOffsets {
		LinkedHashMap<String, Integer> fieldname2offset = new LinkedHashMap<String, Integer>();
		LinkedHashMap<String, VirtualMethodDetails> methodname2offset = new LinkedHashMap<String, VirtualMethodDetails>();

		public ClassOffsets() {
		}

		public ClassOffsets(ClassOffsets parent) {
			fieldname2offset.putAll(parent.fieldname2offset);
			methodname2offset.putAll(parent.methodname2offset);
		}
	}

	HashMap<String, ClassOffsets> classname2offsets = new HashMap<String, TacVisitor.ClassOffsets>();

	void print(String s) {
		sb.append(s);
		sb.append("\n");
	}

	void print_data(String s) {
		sb_data.append(s);
		sb_data.append("\n");
	}

	private void add_string_data(final String lblString, String str) {
		print_data(lblString);
		print_data(Integer.toString(str.length()));
		str = str.replace("\\", "\\\\").replace("\n", "\\n")
				.replace("\t", "\\t").replace("\"", "\\\"");
		print_data("\"" + str + "\"");
	}

	private void add_error_func(String error_postfix) {
		print(error_func_label_prefix + error_postfix);
		print("param " + error_msg_label_prefix + error_postfix);
		print("call " + lblErrorThenExit);
	}

	private void printErrorHandlingFacility() {
		add_string_data(lblMsgNullDereference,
				"Runtime Error: Null pointer dereference!");
		add_string_data(lblMsgArrayOutOfBounds,
				"Runtime Error: Array index out of bounds!");
		add_string_data(lblMsgNegativeArraySize,
				"Runtime Error: Array allocation with negative array size!");
		add_string_data(lblMsgDivByZero, "Runtime Error: Division by zero!");

		print(lblErrorThenExit);
		print("param $0");
		print("call :println");
		print("call :exit");

		add_error_func(null_deref_postfix);
		add_error_func(array_bounds_postfix);
		add_error_func(array_size_postfix);
		add_error_func(div_zero_postfix);
	}

	/*
	 * Convention: Local variables - on current context register, allocated
	 * using next Arrays - allocated per the right size, start of array
	 * allocated using next Label for functions - fullname (class + function
	 * name) every expression is saved in its own register.
	 */

	@Override
	public Object visit(Program program) {
		print("goto " + MAIN_LABEL); // will be set later

		for (DeclClass icClass : program.getClasses()) {
			curr_class = icClass.getName();
			if (icClass.hasSuperClass()) {
				classname2offsets.put(icClass.getName(),
						classname2offsets.get(icClass.getSuperClassName()));
			} else
				classname2offsets.put(icClass.getName(), new ClassOffsets());

			for (DeclField f : icClass.getFields()) {
				f.accept(this);
			}

			for (DeclMethod method : icClass.getMethods()) {
				if (method instanceof DeclStaticMethod)
					continue;

				final String name = method.getName();
				final LinkedHashMap<String, VirtualMethodDetails> offsets_table = classname2offsets
						.get(curr_class).methodname2offset;

				if (offsets_table.containsKey(name)) {
					// Overriding
					offsets_table.get(name).class_name = curr_class;
				} else {
					// Not overriding
					VirtualMethodDetails vmd = new VirtualMethodDetails();
					vmd.class_name = curr_class;
					vmd.method_name = name;
					vmd.index = offsets_table.size();
					offsets_table.put(name, vmd);
				}
			}

			final LinkedHashMap<String, VirtualMethodDetails> method_offsets = classname2offsets
					.get(curr_class).methodname2offset;

			print_data(make_dv_label(curr_class));
			for (String method_name : method_offsets.keySet()) {
				print_data("\t("
						+ make_method_label(
								method_offsets.get(method_name).class_name,
								method_name) + ")");
			}
		}

		for (DeclClass c : program.getClasses()) {
			c.accept(this);
		}

		printErrorHandlingFacility();

		sb.append("\n.data\n");
		sb.append(sb_data);
		return sb.toString();
	}

	@Override
	public Object visit(DeclClass icClass) {
		curr_class = icClass.getName();

		for (DeclMethod m : icClass.getMethods()) {
			m.accept(this);
		}

		return null;
	}

	private static String make_dv_label(String class_name) {
		return ":_" + class_name + "_@DV";
	}

	@Override
	public Object visit(DeclField field) {
		classname2offsets.get(curr_class).fieldname2offset.put(field.getName(),
				classname2offsets.get(curr_class).fieldname2offset.size() + 1);
		return null;
	}

	private String make_method_label(final String class_name,
			final String method_name) {
		return ":" + class_name + "." + method_name;
	}

	private void print_method_code(DeclMethod method) {
		print(make_method_label(curr_class, method.getName()));
		print_method_code_no_label(method);
	}

	private void print_method_code_no_label(DeclMethod method) {
		int taken_by_formals = method.getFormals().size();

		boolean flgSaveRoomForThis = (method instanceof DeclVirtualMethod);
		if (flgSaveRoomForThis) {
			++taken_by_formals;
		}

		context_stack.push(next_register);
		next_register = taken_by_formals;

		for (int i = 0; i < method.getFormals().size(); ++i) {
			method.getFormals().get(i)
					.setRegister(flgSaveRoomForThis ? (i + 1) : i);
		}

		for (Statement s : method.getStatements()) {
			s.accept(this);
		}

		PrimitiveType mock_void = new PrimitiveType(0, DataType.VOID);
		if (method.getType().equals(mock_void)) {
			// add return for the case no one was supplied
			print("ret");
		}

		next_register = context_stack.pop();
	}

	@Override
	public Object visit(DeclVirtualMethod method) {

		print_method_code(method);

		return null;
	}

	@Override
	public Object visit(DeclStaticMethod method) {
		final boolean flgMain = method.getName().equals("main");

		if (flgMain) {
			print(MAIN_LABEL);
			print_method_code_no_label(method);
		} else {
			print_method_code(method);
		}

		return null;
	}

	@Override
	public Object visit(DeclLibraryMethod method) {
		// defined elsewhere
		return null;
	}

	@Override
	public Object visit(Parameter formal) {
		return formal.getRegister();
	}

	@Override
	public Object visit(PrimitiveType type) {
		throw new NotImplementedException();
	}

	@Override
	public Object visit(ClassType type) {
		throw new NotImplementedException();
	}

	private boolean flgNeedConcreteValue = true; // for refArrayElement and
													// refField

	@Override
	public Object visit(StmtAssignment assignment) {
		print("# entered assignment " + assignment.getLine());
		flgNeedConcreteValue = false;
		last_seen_ref_var = null;
		Object res1 = assignment.getVariable().accept(this);
		Node var_seen = last_seen_ref_var;
		last_seen_ref_var = null;
		flgNeedConcreteValue = true;
		Object res2 = assignment.getAssignment().accept(this);
		if (assignment.getVariable() instanceof RefVariable) {
			if (var_seen instanceof DeclField) {
				print("[]= " + res1 + " " + res2 + " #Assignment to field");
			} else {
				print("= " + res2 + " " + res1 + " #Assignment to variable");
			}
		} else if (assignment.getVariable() instanceof RefArrayElement) {
			print("[]= " + res1 + " " + res2 + " #Assignment to array element");
		} else if (assignment.getVariable() instanceof RefField) {
			print("[]= " + res1 + " " + res2 + " #Assignment to field");
		} else {
			throw new NotImplementedException();
		}

		return null;
	}

	@Override
	public Object visit(StmtCall c) {
		return c.getCall().accept(this);
	}

	@Override
	public Object visit(StmtReturn rs) {
		if (rs.hasValue()) {
			Object last_res = rs.getValue().accept(this);
			print("ret " + last_res);
		} else {
			print("ret");
		}
		return null;
	}

	@Override
	public Object visit(StmtIf ifs) {
		final String postfix = "_if_" + Integer.toString(ifs.getLine());
		final String lblAfter = ":after" + postfix;
		final String lblTrue = ":true" + postfix;

		Object res = ifs.getCondition().accept(this);
		print("if " + res + " " + lblTrue);
		if (ifs.hasElse()) {
			ifs.getElseOperation().accept(this);
		}
		print("goto " + lblAfter);

		print(lblTrue);
		ifs.getOperation().accept(this);

		print(lblAfter);

		return null;
	}

	final Stack<String> while_postfix = new Stack<String>();

	private Node last_seen_ref_var;
	final static String label_prefix_after = ":after";
	final static String label_prefix_start = ":start";

	@Override
	public Object visit(StmtWhile ws) {
		final String postfix = "_while_" + Integer.toString(ws.getLine());
		final String lblCurrWhileAfter = label_prefix_after + postfix;
		final String lblCurrWhileStart = label_prefix_start + postfix;

		while_postfix.push(postfix);

		print(lblCurrWhileStart);
		Object cond_res = ws.getCondition().accept(this);
		print("if! " + cond_res + " " + lblCurrWhileAfter);
		ws.getOperation().accept(this);
		print("goto " + lblCurrWhileStart);
		while_postfix.pop();

		print(lblCurrWhileAfter);
		return null;
	}

	@Override
	public Object visit(StmtBreak breakStatement) {
		print("goto " + label_prefix_after + while_postfix.peek());
		return null;
	}

	@Override
	public Object visit(StmtContinue continueStatement) {
		print("goto " + label_prefix_start + while_postfix.peek());
		return null;
	}

	@Override
	public Object visit(StmtBlock stb) {
		for (Statement stmt : stb.getStatements()) {
			stmt.accept(this);
		}
		return null;
	}

	@Override
	public Object visit(LocalVariable lv) {
		lv.setRegister(next_register++);
		if (lv.isInitialized()) {
			Object res = lv.getInitialValue().accept(this);
			print("= " + res + " " + lv.getRegister());
		}
		return lv.getRegister();
	}

	@Override
	public Object visit(RefVariable loc) {
		Node n = loc.getScope().getRecursive(loc.getName());
		last_seen_ref_var = n;

		if (n instanceof DeclField) {
			// we're in a virtual method and need to find the instance base
			// address. It is in $0, by virtual method calling convention
			String class_name = ((DeclClass) n.getScope().owner).getName();
			print("# for DeclField from RefVar " + loc.getName() + " found "
					+ class_name);
			return handle_field_ref(loc.getName(), "$0", class_name);
		} else
			return n.getRegister();
	}

	/*
	 * Both fields and virtual calls refer to instance members. These have 2
	 * steps: 1. Finding which instance. This means finding the location in
	 * memory, resulting in a register containing it. 2. Add offsets to this
	 * base register (for fields) or to the DV (for methods). In any case, this
	 * means *finding the type* of the instance.
	 */

	private String get_base_register(HasObject target) {
		return (String) target.getObject().accept(this);
	}

	private String get_target_class(HasObject nd) {
		return ((ClassType) nd.getObject().getExtraTypeInformation())
				.getClassName();
	}

	private Object handle_field_ref(String field_name, String base_register,
			String class_name) {
		print("# entered field ref " + field_name);
		print("# for field " + field_name + " determined class " + class_name
				+ " and base_register " + base_register);

		int offset = classname2offsets.get(class_name).fieldname2offset
				.get(field_name);
		final String field_val_register = "$" + next_register++;
		print("+ " + base_register + " " + offset + " " + field_val_register
				+ " # for RefField got class_name " + class_name);
		print("# for field ref eval concrete value is " + flgNeedConcreteValue);
		if (flgNeedConcreteValue) {
			print("[] " + field_val_register + " " + field_val_register);
		}
		return field_val_register;
	}

	@Override
	public Object visit(RefField location) {
		String base_register = get_base_register(location);
		String class_name = get_target_class(location);

		return handle_field_ref(location.getField(), base_register, class_name);
	}

	@Override
	public Object visit(VirtualCall call) {
		String base_register;
		String class_name;

		if (!call.hasExplicitObject()) {
			// first, check if static.
			Node n = call.getScope().getRecursive(call.getMethod());
			if (n instanceof DeclStaticMethod) {
				return handleCall(curr_class, call, null, null);
			} else {
				base_register = "$0";
				class_name = ((DeclClass) n.getScope().owner).getName();
			}
		} else {
			base_register = get_base_register(call);
			class_name = get_target_class(call);
		}

		String dv_register = "$" + next_register++;
		print("[] " + base_register + " " + dv_register + " # read DV");

		int offset = classname2offsets.get(class_name).methodname2offset
				.get(call.getMethod()).index;

		print("+ " + dv_register + " " + offset + " " + dv_register
				+ " # for VirtualCall got class_name " + class_name
				+ " and offset " + offset);

		// So now dv_register holds the address to the address of the method
		print("[] " + dv_register + " " + dv_register
				+ " # read method address");

		return handleCall(curr_class, call, dv_register, base_register);
	}

	@Override
	public Object visit(This thisExpression) {
		return "$0";
	}

	@Override
	public Object visit(RefArrayElement loc) {
		boolean flgMyNeedConcreteValue = flgNeedConcreteValue;
		flgNeedConcreteValue = true;

		print("# entered ref array " + loc.getLine());
		print("# " + loc.getLine() + " accessing array on index "
				+ loc.getIndex());
		Object arr = loc.getArray().accept(this);
		gen_null_check(arr);
		Object index = loc.getIndex().accept(this);
		final String address_register = "$" + next_register++;

		gen_bounds_check(arr, index);

		// add 1 to index for skipping the size cell
		final String new_index_reg = "$" + next_register++;
		print("+ " + index + " 1 " + new_index_reg);

		print("+ " + arr + " " + new_index_reg + " " + address_register);
		print("# for array element ref eval concrete value is "
				+ flgMyNeedConcreteValue);
		if (flgMyNeedConcreteValue) {
			print("[] " + address_register + " " + address_register);
		}
		return address_register;
	}

	private void gen_bounds_check(Object arr, Object index) {
		final String len_reg = "$" + next_register++;
		print("[] " + arr + " " + len_reg);
		print(">= " + index + " " + len_reg + " " + len_reg);
		print("if " + len_reg + " " + lblFuncArrayOutOfBounds);
		print("< " + index + " " + 0 + " " + len_reg);
		print("if " + len_reg + " " + lblFuncArrayOutOfBounds);
	}

	private void gen_array_size_check(Object arr_size) {
		final String res_reg = "$" + next_register++;
		print("< " + arr_size + " 0 " + res_reg);
		print("if " + res_reg + " " + lblFuncNegativeArraySize);
	}

	@Override
	public Object visit(StaticCall call) {
		return handleCall(call.getClassName(), call, null, null);
	}

	private Object handleCall(String className, Call call,
			String method_address, String base_address) {
		boolean flgLibraryCall = flgTranslateLibraryCalls
				&& className.equals("Library");

		ArrayList<Object> res_regs = new ArrayList<Object>();

		for (ic.ast.Expression exp : call.getArguments()) {
			res_regs.add(exp.accept(this));
		}

		if (!flgLibraryCall) {
			if (base_address != null) {
				print("# print base_address or 0 for static");
				print("param " + base_address);
			} else {
				// print("param 0");
			}
		}

		for (Object res : res_regs) {
			print("param " + res);
		}

		if (method_address == null) {
			String lblTargetFunction;
			if (flgLibraryCall) {
				lblTargetFunction = ":" + call.getMethod();
			} else {
				lblTargetFunction = ":" + className + "." + call.getMethod();
			}
			method_address = lblTargetFunction;
		}

		PrimitiveType mock_void = new PrimitiveType(0, DataType.VOID);

		if (!call.getExtraTypeInformation().equals(mock_void)) {
			String call_result_register = "$" + next_register++;
			print("call " + method_address + " " + call_result_register);
			return call_result_register;
		} else {
			print("call " + method_address);
			return null;
		}
	}

	@Override
	public Object visit(NewInstance nc) {
		ClassOffsets co = classname2offsets.get(nc.getName());
		int size = co.fieldname2offset.size() + 1; // 1 for dv + 1 for each
													// field
		print("param " + size);
		String obj_address_reg = "$" + next_register++;
		print("call :alloc " + obj_address_reg);
		print("[]= " + obj_address_reg + " " + make_dv_label(nc.getName())
				+ " #write DV to head of alloc");

		return obj_address_reg;
	}

	@Override
	public Object visit(NewArray na) {
		Object arr_size = na.getSize().accept(this);
		gen_array_size_check(arr_size);
		print("# begin allocating new array of size " + arr_size);

		// The following commented out code is wrong as allocations of new
		// arrays is of the form int x[][] = new T[N]
		// if (na.getType().getArrayDimension() > 1) {
		// String new_size_register = "$" + next_register++;
		// print("* " + arr_size + " " + na.getType().getArrayDimension()
		// + " " + new_size_register);
		// arr_size = new_size_register;
		// }

		final String full_size_register = "$" + next_register++; // size to be
																	// allocated,
																	// including
																	// room for
																	// size
																	// itself
		print("+ " + arr_size + " 1 " + full_size_register);

		String arr_address_reg = "$" + next_register++;
		print("param " + full_size_register);
		print("call :alloc " + arr_address_reg);

		// Apparently Shachar says in forum all arrays should have size in first
		// cell -> easier to just handle it like it later, to support library
		// calls
		//

		print("[]= " + arr_address_reg + " " + arr_size
				+ " #write size to head of alloc");

		return arr_address_reg;
	}

	@Override
	public Object visit(Length len) {
		print("# starting to get array for length:");
		Object arr = len.getArray().accept(this);
		gen_null_check(arr);
		print("# ended getting array for length");
		String res_reg = "$" + next_register++;
		print("[] " + arr + " " + res_reg + " #length");

		return res_reg;
	}

	private void gen_null_check(Object o) {
		print("if! " + o + " " + lblFuncNullDereference);
	}

	private int next_string_num = 0;

	@Override
	public Object visit(Literal literal) {
		String val_to_store;
		switch (literal.getType()) {
		case BOOLEAN:
			val_to_store = literal.getValue().equals("true") ? "1" : "0";
			break;
		case INT:
			val_to_store = literal.getValue().toString();
			break;
		case STRING:
			final String str = literal.getValue().toString();
			final String lblString = ":string_" + next_string_num++;
			add_string_data(lblString, str);
			val_to_store = lblString;
			break;
		case VOID:
			// null
			return 0;
		default:
			throw new NotImplementedException();
		}

		String register = "$" + next_register++;
		print("= " + val_to_store + " " + register);
		return register;
	}

	@Override
	public Object visit(UnaryOp op) {
		Object opd = op.getOperand().accept(this);
		final String res_reg = "$" + next_register++;

		switch (op.getOperator()) {
		case LNEG:
			print("! " + opd + " " + res_reg);
			break;
		case UMINUS:
			print("- 0 " + opd + " " + res_reg);
			break;
		default:
			throw new NotImplementedException();
		}

		return res_reg;
	}

	private int curr_op_line = -1;
	private int next_op_in_line_num = 0;

	@Override
	public Object visit(BinaryOp op) {
		if (curr_op_line != op.getLine()) {
			next_op_in_line_num = 0;
			curr_op_line = op.getLine();
		}

		Object opd1 = op.getFirstOperand().accept(this);
		final String res_reg = "$" + next_register++;

		final String lblAfterEvalAll = ":after_eval_all_" + op.getLine() + "_"
				+ next_op_in_line_num++;

		if (op.getOperator() == BinaryOps.LAND
				|| op.getOperator() == BinaryOps.LOR) {
			print("&& " + opd1 + " 1 " + res_reg);
			String stop_if_first = (op.getOperator() == BinaryOps.LAND ? "if!"
					: "if");
			print(stop_if_first + " " + opd1 + " " + lblAfterEvalAll);
		}

		Object opd2 = op.getSecondOperand().accept(this);

		String int_op = null;

		switch (op.getOperator()) {
		case DIVIDE:
			if (int_op == null) {
				int_op = "/";
				gen_check_non_zero(opd2);
			}
		case GT:
			if (int_op == null)
				int_op = ">";
		case GTE:
			if (int_op == null)
				int_op = ">=";
		case LT:
			if (int_op == null)
				int_op = "<";
		case LTE:
			if (int_op == null)
				int_op = "<=";
		case MINUS:
			if (int_op == null)
				int_op = "-";
		case MOD:
			if (int_op == null) {
				int_op = "%";
				gen_check_non_zero(opd2);
			}
		case MULTIPLY:
			if (int_op == null)
				int_op = "*";
		case LAND:
			if (int_op == null)
				int_op = "&&";
		case LOR:
			if (int_op == null)
				int_op = "||";
		case EQUAL:
			if (int_op == null)
				int_op = "==";
		case NEQUAL:
			if (int_op == null)
				int_op = "!=";
		case PLUS:
			if (int_op == null)
				int_op = "+";
			break;
		}

		PrimitiveType mock_string = new PrimitiveType(0, DataType.STRING);
		if (op.getOperator().equals(BinaryOp.BinaryOps.PLUS)
				&& op.getFirstOperand().getExtraTypeInformation()
						.equals(mock_string)) {
			print("param " + opd1);
			print("param " + opd2);
			print("call :stringCat " + res_reg);
		} else {
			print(int_op + " " + opd1 + " " + opd2 + " " + res_reg);
		}

		print(lblAfterEvalAll);

		return res_reg;
	}

	private void gen_check_non_zero(Object opd) {
		print("if! " + opd + " " + lblFuncDivByZero);
	}
}
